#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
更新现有conda环境中的包，基于requirements.txt文件
"""

import os
import sys
import subprocess
import argparse
from pathlib import Path
import tempfile

# 常见的conda专用包，这些包通常不能通过pip安装
CONDA_SPECIFIC_PACKAGES = [
    'bzip2', 'ca-certificates', 'libarchive', 'libcurl', 'libffi', 'libiconv',
    'libmamba', 'libsolv', 'libssh2', 'libxml2', 'lz4-c', 'openssl', 'pcre2',
    'tk', 'xz', 'yaml-cpp', 'zlib', 'zstandard', 'zstd', 'vc', 'vs2015_runtime'
]

def list_available_requirements():
    """列出可用的requirements.txt文件"""
    env_dir = Path("conda_environments")
    if not env_dir.exists():
        print(f"错误: 找不到环境目录 {env_dir}")
        return []
    
    req_files = list(env_dir.glob("*_requirements.txt"))
    return [f.name for f in req_files]

def split_requirements(req_file):
    """将requirements文件分成conda包和pip包"""
    conda_packages = []
    pip_packages = []
    
    with open(req_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
                
            # 检查是否是conda专用包
            is_conda_package = False
            for pkg in CONDA_SPECIFIC_PACKAGES:
                if line.startswith(pkg + "==") or line == pkg:
                    conda_packages.append(line)
                    is_conda_package = True
                    break
            
            if not is_conda_package:
                pip_packages.append(line)
    
    return conda_packages, pip_packages

def update_environment(req_file, env_name, use_tsinghua=True):
    """更新conda环境中的包"""
    req_path = Path("conda_environments") / req_file
    
    if not req_path.exists():
        print(f"错误: requirements文件 {req_file} 不存在")
        return False
    
    print(f"正在使用文件 {req_file} 更新conda环境 {env_name}...")
    
    # 检查环境是否存在
    result = subprocess.run(["conda", "env", "list"], capture_output=True, text=True)
    if env_name not in result.stdout:
        print(f"错误: 环境 {env_name} 不存在")
        return False
    
    # 设置环境变量
    env = os.environ.copy()
    
    # 添加清华镜像源参数
    if use_tsinghua:
        # 设置conda配置使用清华镜像源
        subprocess.run(["conda", "config", "--set", "show_channel_urls", "yes"], 
                    check=False)
        subprocess.run(["conda", "config", "--add", "channels", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge"], 
                    check=False)
        subprocess.run(["conda", "config", "--add", "channels", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free"], 
                    check=False)
        subprocess.run(["conda", "config", "--add", "channels", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main"], 
                    check=False)
        
        print("已配置清华镜像源")
    
    # 分割requirements文件
    conda_packages, pip_packages = split_requirements(req_path)
    
    # 使用conda更新conda专用包
    if conda_packages:
        print(f"正在使用conda更新环境 {env_name} 中的conda专用包...")
        
        # 创建临时文件存储conda包
        with tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix='.txt') as tmp:
            for pkg in conda_packages:
                tmp.write(f"{pkg}\n")
            tmp_path = tmp.name
        
        try:
            # 使用conda安装包
            process = subprocess.run(["conda", "install", "--name", env_name, "--file", tmp_path, "--yes"], 
                                  env=env)
            
            if process.returncode != 0:
                print(f"conda包更新失败，返回码: {process.returncode}")
                # 继续尝试安装pip包
        finally:
            # 删除临时文件
            os.unlink(tmp_path)
    
    # 使用pip更新其他包
    if pip_packages:
        print(f"正在使用pip更新环境 {env_name} 中的pip包...")
        
        # 创建临时文件存储pip包
        with tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix='.txt') as tmp:
            for pkg in pip_packages:
                tmp.write(f"{pkg}\n")
            tmp_path = tmp.name
        
        try:
            # 使用pip安装包
            process = subprocess.run(["conda", "run", "-n", env_name, "pip", "install", "-r", tmp_path, "--upgrade"], 
                                   env=env)
            
            if process.returncode != 0:
                print(f"pip包更新失败，返回码: {process.returncode}")
                return False
        finally:
            # 删除临时文件
            os.unlink(tmp_path)
    
    print("环境更新成功!")
    return True

def main():
    parser = argparse.ArgumentParser(description="更新现有conda环境中的包，基于requirements.txt文件")
    parser.add_argument("--req-file", type=str, help="requirements.txt文件名 (位于conda_environments目录下)")
    parser.add_argument("--env-name", type=str, help="要更新的conda环境名称")
    parser.add_argument("--list", action="store_true", help="列出可用的requirements.txt文件")
    parser.add_argument("--no-tsinghua", action="store_true", help="不使用清华镜像源")
    
    args = parser.parse_args()
    
    if args.list:
        print("可用的requirements.txt文件:")
        for req in list_available_requirements():
            print(f"- {req}")
        return
    
    # 检查是否提供了环境名称
    if not args.env_name:
        print("错误: 请提供要更新的环境名称 (--env-name)")
        return
    
    if not args.req_file:
        # 如果没有指定requirements文件，列出可用选项并让用户选择
        available_reqs = list_available_requirements()
        if not available_reqs:
            return
        
        print("请选择要使用的requirements.txt文件:")
        for i, req in enumerate(available_reqs):
            print(f"{i+1}. {req}")
        
        try:
            choice = int(input("请输入选项编号: "))
            if 1 <= choice <= len(available_reqs):
                req_file = available_reqs[choice-1]
            else:
                print("无效的选择")
                return
        except ValueError:
            print("请输入有效的数字")
            return
    else:
        req_file = args.req_file
    
    update_environment(req_file, args.env_name, not args.no_tsinghua)

if __name__ == "__main__":
    main() 